from langchain_community.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableSequence
import os

# 初始化OpenAI语言模型
llm = ChatOpenAI(model_name="gpt-4", base_url=os.environ["OPENAI_BASE_URL"])

# 定义一个简单的提示模板
template = "您是一个乐于助人的助手，请回答一下问题:{question}"

# PromptTemplate适用于用于生成单一的字符串提示，不能用于多角色的聊天
prompt = PromptTemplate(template=template, input_variables=["question"])

# #  解析结果 用过滤chatModel的输出
parser = StrOutputParser()
# 创建一个LLMChain实例
# chain = RunnableSequence(prompt | llm | parser)
chain = prompt | llm | parser

# 使用chain来生成一个回答
question = input("请输入你的问题：")
response = chain.invoke({"question": question})
print(f"Ai大模型的回答: {response}")